from functools import wraps
from typing import Any, Callable, Literal, Optional, Type, Union

from sqlalchemy import create_engine, Engine, URL
from sqlalchemy.orm import Session, sessionmaker, SessionTransaction

from ._logger import logger

__all__ = [
    'Engine',
    'Session',
    'check_engine_is_initialized',
    'get_db',
    'create_session',
    'translation',
    'build_url',
    'initialize',
    'initialize_for_production',
    'initialize_for_developmental',
]

Engine = Engine
Session = Session
# noinspection SpellCheckingInspection
SQLURL = Union[str, URL]

# 全局数据库引擎
engine: Optional[Engine] = None
# 全局 Session 构造器
session: Optional[Type[Session]] = None


def check_engine_is_initialized():
    """
    检查数据库引擎是否初始化，若未初始化，则抛出异常
    """
    if engine is None:
        raise RuntimeError('Please initialize the database instance first.')


def get_db() -> Engine:
    """
    获取数据库引擎对象
    """
    check_engine_is_initialized()
    return engine


def create_session() -> Session:
    """
    创建会话
    """
    check_engine_is_initialized()
    return session()


def translation(func: Callable[[Session, Any], Any]):
    """
    SQL 事务拦截器

    :param func: 被监听的方法
    :return: `func` 的返回值
    """

    @wraps(func)
    def wrapper(*args, **kwargs) -> Any:
        tx: Optional[SessionTransaction] = None
        _session = create_session()
        try:
            # 开启事务
            tx = _session.begin()
            # 调用被监听的函数
            result = func(_session, *args, **kwargs)
            # 提交事务
            tx.commit()
            # 返回结果
            return result
        except BaseException as err:
            # 若发生异常，则回滚事务
            if tx is not None:
                tx.rollback()
            logger.error(err)
            raise err
        finally:
            # 关闭事务与数据库会话
            if tx is not None:
                tx.close()
            if _session is not None:
                _session.close()

    return wrapper


def build_url(
        driver: Literal["sqlite", "mysql", "mariadb", "postgresql"],
        host: str = None, port: int = None, database: str = None,
        username: str = None, password: str = None,
) -> SQLURL:
    """
    创建数据库连接地址

    :param driver: 驱动名称
    :param host: 数据库主机名
    :param port: 数据库端口号
    :param database: 数据库名称
    :param username: 用户名
    :param password: 密码
    :return: 数据库连接地址
    """
    return URL.create(
        drivername=driver, host=host, port=port, database=database,
        username=username, password=password,
    )


def initialize(url: SQLURL, db_configs=None, session_configs=None):
    """
    初始化数据库

    :param url: 数据库连接地址
    :param db_configs: 数据库配置
    :param session_configs: 数据库会话配置
    """
    # 获取全局引用
    global engine, session
    if engine is not None:
        # 若数据库引擎不为 None，则发出报警
        logger.warning('The database has been initialized and cannot be initialized again.')
        return
    # 创建数据库实例
    engine = create_engine(url=url, **(db_configs if db_configs is not None else {}))
    # 创建会话实例
    session = sessionmaker(bind=engine, **(session_configs if session_configs is not None else {}))


def initialize_for_production(url: SQLURL, db_configs=None, session_configs=None):
    """
    初始化用于生产环境的数据库实例

    :param url: 数据库连接地址
    :param db_configs: 数据库配置
    :param session_configs: 数据库会话配置
    """
    if db_configs is None:
        db_configs = {}
    db_configs.update(
        {
            'echo': False,
            'echo_pool': False,
        }
    )
    initialize(
        url,
        db_configs,
        session_configs,
    )


def initialize_for_developmental(url: SQLURL, db_configs=None, session_configs=None):
    """
    初始化用于开发环境的数据库实例

    :param url: 数据库连接地址
    :param db_configs: 数据库配置
    :param session_configs: 数据库会话配置
    """
    if db_configs is None:
        db_configs = {}
    db_configs.update(
        {
            'echo': True,
            'echo_pool': True,
        }
    )
    initialize(
        url,
        db_configs,
        session_configs,
    )
